#include <iostream>
using namespace std;
#include <unistd.h>
#include <string>
#include <vector>
#include <stdlib.h>
#include <time.h>
#include <sys/wait.h>
#include "tasks.hpp"

enum
{
    UsageError = 1,
    ArgError,
    PipeError
};

void Usage(char *argv)
{
    cout << "Usage:" << endl;
    cout << "\t" << argv << " sub_process_num" << endl;
}

class Channel
{
public:
    Channel(int wfd, pid_t sub_process_id, const string name)
        : _wfd(wfd), _sub_process_id(sub_process_id), _name(name)
    {
    }

    void Debug()
    {
        cout << "_wfd: " << _wfd;
        cout << ", _sub_process_id: " << _sub_process_id;
        cout << ", _name: " << _name << endl;
    }

    void Close()
    {
        close(_wfd);
    }

    int Wfd() { return _wfd; }

    string Name() { return _name; }

    pid_t Pid() { return _sub_process_id; }

    ~Channel() {}

private:
    int _wfd;
    pid_t _sub_process_id;
    string _name;
};

class ProcessPool
{
public:
    ProcessPool(int sub_process_num)
        : _sub_process_num(sub_process_num)
    {
    }

    // *** vector<Channel>&
    int CreatChannels(work_t work)
    {
        vector<int> fds;
        // 1. 创建通信信道和进程
        for (int i = 0; i < _sub_process_num; i++)
        {
            // 先创建管道
            int pipefd[2];
            int n = pipe(pipefd);
            if (n < 0)
            {
                return PipeError;
            }
            
            fds.push_back(pipefd[1]); // 保存管道的写端

            // 再创建子进程，确保父进程和子进程
            pid_t id = fork();
            if (id == 0)
            {
                // child -> r
                //close(pipefd[1]);
                if (!fds.empty())
                {
                    for (auto& fd : fds)
                    {
                        close(fd);
                    }
                }

                // TODO
                dup2(pipefd[0], 0);
                work(pipefd[0]);

                // sleep(100);
                exit(0);
            }

            // father
            close(pipefd[0]);
            string cname = "channel--" + to_string(i);

            _channels.push_back(Channel(pipefd[1], id, cname));
            // fds.push_back(pipefd[1]); // 保存管道的写端
        }

        return 0;
    }

    void Debug()
    {
        for (auto e : _channels)
        {
            e.Debug();
        }
    }

    int NextChannel()
    {
        static int next = 0;
        int c = next;
        ++next;
        next %= _channels.size();

        return c;
    }

    void SendTask(int index, int command_code)
    {
        cout << "Send task to " << _channels[index].Name() << ", sub process id: " << _channels[index].Pid() << endl;
        write(_channels[index].Wfd(), &command_code, sizeof(command_code));
    }

    void KillAll()
    {
        for (auto& channel : _channels)
        {
            close(channel.Wfd());

            pid_t pid = channel.Pid(); // 子进程(管道读端进程)的 pid   
            pid_t rid = waitpid(pid, nullptr, 0);
            if (rid == pid) // wait 成功
            {
                cout << "wait sub process: " << pid << "success..." << endl;
            }
            cout << "close channel: " << channel.Name() << ", sub process is to quit.." << endl;
        }
    }

    ~ProcessPool()
    {
    }

private:
    vector<Channel> _channels;
    int _sub_process_num;
};

void CtrlProcess(ProcessPool *ptr, int cnt)
{
    while (cnt)
    {
        // a. 选择一个通道和进程
        int channel = ptr->NextChannel();
        // b. 选择一个任务
        int task = NextTask();
        // c. 发送任务
        ptr->SendTask(channel, task);

        sleep(1); // 每隔 1s 发送一次任务
        --cnt;
    }
}

int main(int argc, char *argv[])
{
    if (argc < 2)
    {
        Usage(argv[0]);
        return UsageError;
    }

    int sub_process_num = stoi(argv[1]);
    if (sub_process_num <= 0)
    {
        return ArgError;
    }

    // 1. 创建通信信道和进程
    ProcessPool *processpool_ptr = new ProcessPool(sub_process_num);
    processpool_ptr->CreatChannels(worker);
    // processpool_ptr->Debug();

    // sleep(100);
    srand(time(nullptr));

    // 2. 控制子进程

    CtrlProcess(processpool_ptr, 5); // 假定 10 个任务
    cout << "task run done" << endl;

    // for (auto channel : channels)
    // {
    //     channel.Debug();
    // }

    // 3. 回收进程
    processpool_ptr->KillAll();

    // sleep(100);
    delete processpool_ptr;
    return 0;
}